"""
MEMESA-TOOLS
============
Model preparation, vertex enumeration and vertex analysis tools developed as
part of the MEMESA project.

    Copyright (C) 2009-2012 Brett G. Olivier

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>

Author: Brett G. Olivier
Contact email: bgoli@users.sourceforge.net

"""

import os, time
cDir = os.path.dirname(os.path.abspath(os.sys.argv[0]))
ALTLOADFILE = None

# Set userdata here #

# yeast 5.32
#work_dir = 'yeast532'
#model_name = 'yeast_5.33_model' # base name not xml file

# big Ecoli
work_dir = 'iAF1260'
model_name = 'Ecoli_iAF1260_noox.glc' # base name not xml file
##  ALTLOADFILE = 'Ecoli_iAF1260_ox.glc.None.hdf5'

# small Ecoli
# work_dir = 'iJR904'
# model_name = 'Ecoli_iJR904.ac' # base name not xml file
# model_name = 'Ecoli_iJR904.glc' # base name not xml file
# model_name = 'Ecoli_iJR904.glc2' # base name not xml file
# model_name = 'Ecoli_iJR904.lcts' # base name not xml file
# model_name = 'Ecoli_iJR904.succ' # base name not xml file

# ALTLOAD do not do .ine to hdf5 conversion simply load the specified hdf5 file
ALTLOAD = False
# use the *.columns.txt file from the Steven output to cross check flux order
USE_COLUMN_CROSSCHECK = True
# rounding off digits
rFact = 15
# write vertex_variable_flux.csv
FIND_VARIABLE_FLUXES_IN_VERTICES = True # needed for module analysis

#  write vertex_array.csv
WRITE_VERTEX_ARRAY = False # not needed
WRITE_VERTEX_ARRAY_NEW = False # for SMALL MODELS ONLY
WRITE_VERTEX_ARRAY_FULL = False # once only small models

# memory efficient algorithms
BIG_FILE = True

# end userdata here #

model_dir = os.path.join(cDir, work_dir, 'sbml')
work_dir = os.path.join(cDir, work_dir, 'vertex')
model_file = '%s.xml' % model_name
vertex_file = os.path.join(work_dir, '%s.noinf_r.ine.all' % model_name)
rfva_file = os.path.join(work_dir, '%s.noinf_r.ine.opt.fva' % model_name)

### Active script ###

if ALTLOAD:
    ALTLOADFILE = os.path.join(work_dir, ALTLOADFILE)
    assert os.path.exists(ALTLOADFILE)

import numpy
import sympy
import h5py
import gc
gc.enable()

from pyscescbm.CBVersion import __DEBUG__, __version__
from pyscescbm import CBRead, CBWrite, CBTools
from pyscescbm import CBSolver as slv

if not os.path.exists(work_dir):
    os.mkdir(work_dir)

mod = CBRead.readSBML2FBA(model_file, work_dir=model_dir)
mod.id = model_name
CBTools.addStoichToFBAModel(mod)
CBTools.processBiGGchemFormula(mod)
CBTools.processBiGGannotationNote(mod, annotation_key='note')

if USE_COLUMN_CROSSCHECK:
    mismatch = []
    F = file(os.path.join(work_dir, '%s.noinf_r.columns.txt' % model_name), 'r')
    indx_lst = []
    for L in F:
        l = [i.strip() for i in L.split(',')]
        if len(l) == 2:
            indx_lst.append(l[1])
    F.close()
    for r in range(len(indx_lst)):
        match = bool(indx_lst[r] == mod.N.col[r])
        print '%s == %s: %s' % (indx_lst[r], mod.N.col[r], match)
        if not match: mismatch.append((indx_lst[r], mod.N.col[r]))
    if len(mismatch) > 0:
        print '\nColumn mismatch (%s)\n' % len(mismatch)
        ##  print mismatch
        os.sys.exit(1)
    else:
        print '\nColumn check successful, have a nice day :-)'
        time.sleep(2)


TIME_START = time.time()
# get Real FVA results from Steven
realFVAdata = CBRead.readSK_FVA(rfva_file)

for f in range(len(realFVAdata)):
    realFVAdata[f].update({'id' : mod.N.col[f]})

for j in realFVAdata:
    print j

realNamesVar = [j['id'] for j in realFVAdata if j['status']=='VARIABLE']
realFVA_vari_len = len(realNamesVar)
realNamesFix = [j['id'] for j in realFVAdata if j['status']=='FIXED']
realFVA_fixed_len = len(realNamesFix)
print "\nReal FVA reports %s fixed and %s variable fluxes." % (realFVA_fixed_len, realFVA_vari_len)
print realFVAdata[:5]

FLUX_NAMES = tuple(mod.N.col)

# parse vertex file to array and symb_array

if not ALTLOAD:
    HD5datF = CBRead.readSK_vertex(vertex_file, bigfile=True, fast_rational=True, nformat='%.15f', compression=None, hdf5file=os.path.join(work_dir, model_name))
else:
    HD5datF = ALTLOADFILE
HD5dat = h5py.File(HD5datF,'r')

print '\nUsing HDF5 datafile: %s\n' % HD5datF

HAVE_RAYS = True
HAVE_LIN = True

Tlist = []
try:
    for l in range(HD5dat['data/rays'].shape[0]):
        Tlist.append(HD5dat['data/rays'][l])
    rayList = Tlist
except:
    print 'No Rays'
    rayList = []
    HAVE_RAYS = False
Tlist = []
try:
    for l in range(HD5dat['data/lin'].shape[0]):
        Tlist.append(HD5dat['data/lin'][l])
    linBasis = Tlist
except:
    print 'No Lineality vectors'
    linBasis = []
    HAVE_LIN = False
del Tlist
print 'Rays & Linealities', len(rayList), len(linBasis)

rayF = file(os.path.join(work_dir,'%s.rays.csv' % model_name),'w')
if HAVE_RAYS:
    # find fluxes active in rays
    ray_arr = numpy.array(rayList)
    RAY_FLUXES = []
    del rayList
    print ray_arr.shape
    for c in range(ray_arr.shape[1]):
        if (ray_arr[:,c] > 0.0).any():
            print FLUX_NAMES[c], (ray_arr[:,c] > 0.0).any()
            RAY_FLUXES.append(FLUX_NAMES[c])

    print len(RAY_FLUXES)
    print RAY_FLUXES

    # write the rays
    for c in range(len(realFVAdata)):
        if c != len(realFVAdata)-1:
            rayF.write('%s,' % (c+1))
        else:
            rayF.write('%s\n' % (c+1))
    for c in range(len(realFVAdata)):
        id = realFVAdata[c]['id']
        if c != len(realFVAdata)-1:
            rayF.write('%s,' % id)
        else:
            rayF.write('%s\n' % id)
    ray_arr.shape
    for r in range(ray_arr.shape[0]):
        for c in range(ray_arr.shape[1]):
            if c != ray_arr.shape[1]-1:
                rayF.write('%s,' % ray_arr[r,c])
            else:
                rayF.write('%s\n' % ray_arr[r,c])
    del ray_arr
rayF.flush()
rayF.close()

linbF = file(os.path.join(work_dir,'%s.lineality_basis.csv' % model_name),'w')
if HAVE_LIN:
    # find fluxes active in lineality space
    linbasis_arr = numpy.array(linBasis)
    del linBasis
    print linbasis_arr.shape

    LIN_FLUXES = []
    for c in range(linbasis_arr.shape[1]):
        if (linbasis_arr[:,c] > 0.0).any():
            print FLUX_NAMES[c], (linbasis_arr[:,c] > 0.0).any()
            LIN_FLUXES.append(FLUX_NAMES[c])

    print len(LIN_FLUXES)
    print LIN_FLUXES

    # write the basis for the lineality space
    for c in range(len(realFVAdata)):
        if c != len(realFVAdata)-1:
            linbF.write('%s,' % (c+1))
        else:
            linbF.write('%s\n' % (c+1))
    for c in range(len(realFVAdata)):
        id = realFVAdata[c]['id']
        if c != len(realFVAdata)-1:
            linbF.write('%s,' % id)
        else:
            linbF.write('%s\n' % id)
    for r in range(linbasis_arr.shape[0]):
        for c in range(linbasis_arr.shape[1]):
            if c != linbasis_arr.shape[1]-1:
                linbF.write('%s,' % linbasis_arr[r,c])
            else:
                linbF.write('%s\n' % linbasis_arr[r,c])
    del linbasis_arr
linbF.flush()
linbF.close()


# get the vertices as a list/array
vertList = HD5dat['data/vertices']
print len(vertList)

# populate matrices with fixed data
vertex_arr_status = []
vertex_arr_min = []
vertex_arr_max = []
for j in realFVAdata:
    j_min = sympy.Rational('%s' % j['min'])
    j_max = sympy.Rational('%s' % j['max'])
    vStat = None
    if j['status']=='FIXED':
        vStat = 1
    else:
        vStat = 0
    vertex_arr_status.append(vStat)
    vertex_arr_min.append(j_min.evalf())
    vertex_arr_max.append(j_max.evalf())
del j_min, j_max

"""
create a map index file for vertex_array with the following columns:

 - index
 - id
 - name
 - subsystem
 - fixed (1/0)
 - FVA min
 - FVA max

The last three columns are included in vertex_array as the first 3 rows
"""
# write vertex_array_header.csv
mapF = file(os.path.join(work_dir,'%s.vertex_array_header.csv' % model_name),'w')
mapF.write('Index,Id,Name,Subsystem,Fixed,FVAmin,FVAmax,Equation\n')
for r in range(len(realFVAdata)):
    R = mod.getReaction(realFVAdata[r]['id'])
    name = '"%s"' % R.name
    if R.annotation.has_key('SUBSYSTEM'):
        subsystem = '"%s"' % R.annotation['SUBSYSTEM']
    else:
        subsystem = ''
    if R.annotation.has_key('Equation'):
        equation = '"%s"' % R.annotation['Equation']
    else:
        equation = ''
    mapF.write('%s,%s,%s,%s,%i,%s,%s,%s\n' % (r+1, realFVAdata[r]['id'], name, subsystem, vertex_arr_status[r], vertex_arr_min[r], vertex_arr_max[r], equation))
mapF.close()
gc.collect()

#  write vertex_array_full.csv
if WRITE_VERTEX_ARRAY_FULL:
    CBTools.exportLabelledLinkedList([vertex_arr_status], None, os.path.join(work_dir,'%s.vertex_array_full.csv' % model_name), sep=',', format='%s', appendlist=False)
    CBTools.exportLabelledLinkedList([vertex_arr_min], None, os.path.join(work_dir,'%s.vertex_array_full.csv' % model_name), sep=',', format='%s', appendlist=True)
    CBTools.exportLabelledLinkedList([vertex_arr_max], None, os.path.join(work_dir,'%s.vertex_array_full.csv' % model_name), sep=',', format='%s', appendlist=True)

    # write out chunks
    cntr = 0
    cntr2 = 2000
    while cntr < len(vertList)-50 and cntr2 < len(vertList)-1:
        CBTools.exportLabelledLinkedList(vertList[cntr:cntr2], None, os.path.join(work_dir,'%s.vertex_array_full.csv' % model_name), sep=',', format='%s', appendlist=True)
        cntr = cntr2
        cntr2 += 2000
        print cntr,
    print ' '
    # write out the remaining bit
    for x in xrange(cntr, len(vertList)):
        CBTools.exportLabelledLinkedList([vertList[x]], None, os.path.join(work_dir,'%s.vertex_array_full.csv' % model_name), sep=',', format='%s', appendlist=True)
    print x+1


#  write vertex_array.csv
def writeVertexArray(rng, model_name, mod, vertList, vertex_arr_status, vertex_arr_min, vertex_arr_max):
    fname = os.path.join(work_dir,'%s.vertex_array.csv' % (model_name))
    FO = file(fname, 'w')
    cntr = 1
    fcntr = 0
    col_l = len(vertex_arr_status)
    row_l = vertList.shape[0]
    print 'vl', row_l, col_l, vertList.shape[1], rng
    Btime = time.time()
    for col in xrange(*rng):
        if vertex_arr_status[col] == 0.0:
            print time.strftime('%d-%H:%M'), '%3.2f' % (float(cntr)/float(realFVA_vari_len)*100), '%2.2f' % (float(time.time()-Btime)/60.0), mod.N.col[col], col
            cntr += 1
            R = mod.getReaction(mod.N.col[col])
            if R.annotation.has_key('SUBSYSTEM'):
                subsystem = '"%s"' % R.annotation['SUBSYSTEM']
            else:
                subsystem = ''
            if R.annotation.has_key('Equation'):
                equation = '"%s"' % R.annotation['Equation']
            else:
                equation = ''
            FO.write('"%s","%s",%s,%s,%s,%s,' % (R.getPid(), R.name, equation, subsystem, vertex_arr_min[col],vertex_arr_max[col]))
            outStr = ''
            rcntr = 0

            for row in xrange(0, row_l):
                rcntr += 1
                if row != row_l-1:
                    if vertList[row][col] == 0.0:
                        outStr += '0.0,'
                    else:
                        outStr += '%s,' % round(vertList[row][col],rFact)
                else:
                    if vertList[row][col] == 0.0:
                        outStr += '0.0\n'
                    else:
                        outStr += '%s\n' % round(vertList[row][col],rFact)
                if rcntr == 5000:
                    FO.write(outStr)
                    rcntr = 0
                    outStr = ''
            FO.write(outStr)
            FO.flush()
            outStr = ''
    FO.close()
    return fname

def writeVertexArray_new(rng, model_name, mod, vertList, var_fluxes, output_array=False):
    fname = os.path.join(work_dir,'%s.vertex_array.NEW.csv' % (model_name))
    FO = file(fname, 'w')
    cntr = 1
    fcntr = 0
    col_l = len(vertex_arr_status)
    row_l = vertList.shape[0]
    print 'vl', row_l, col_l, vertList.shape[1], rng
    Btime = time.time()
    varFluxesL = len(var_fluxes)
    print 'Variable fluxes:', varFluxesL
    Btime = time.time()
    fcntr = 0
    for col in xrange(*rng):
        if mod.N.col[col] in var_fluxes:
            print time.strftime('%d-%H:%M'), '%3.2f' % (float(cntr)/float(realFVA_vari_len)*100), '%2.2f' % (float(time.time()-Btime)/60.0), mod.N.col[col], col
            cntr += 1
            R = mod.getReaction(mod.N.col[col])
            FO.write('%s' % R.getPid())
            if fcntr != varFluxesL-1:
                FO.write(',')
            else:
                FO.write('\n')
            fcntr += 1
    Btime = time.time()
    OUTPUT_ARRAY = []
    cntr = 0
    for row in xrange(0, row_l):
        if cntr == 0 or cntr == 19:
            print time.strftime('%d-%H:%M'), '%3.2f' % (float(row)/float(row_l)*100), '%2.2f' % (float(time.time()-Btime)/60.0), row
            cntr = 0
        cntr += 1
        outStr = ''
        fcntr = 0
        OUTPUT_COL = []
        for col in xrange(*rng):
            if mod.N.col[col] in var_fluxes:
                if output_array:
                    OUTPUT_COL.append(vertList[row][col])
                if fcntr != varFluxesL-1:
                    if vertList[row][col] == 0.0:
                        outStr += '0.0,'
                    else:
                        outStr += '%s,' % round(vertList[row][col],rFact)
                else:
                    if vertList[row][col] == 0.0:
                        outStr += '0.0\n'
                    else:
                        outStr += '%s\n' % round(vertList[row][col],rFact)
                fcntr += 1
        if output_array:
            OUTPUT_ARRAY.append(OUTPUT_COL)
        FO.write(outStr)
        FO.flush()
        outStr = ''

    FO.close()
    if output_array:
        return fname, numpy.array(OUTPUT_ARRAY)
    else:
        return fname, None

#  write vertex_array.csv
if WRITE_VERTEX_ARRAY:
    fname = writeVertexArray((0,len(vertex_arr_status)), model_name, mod, vertList, vertex_arr_status, vertex_arr_min, vertex_arr_max)
# find constant fluxes in vertices export as (write vertex_variable_fluxes.csv)
if FIND_VARIABLE_FLUXES_IN_VERTICES:
    vrl = len(vertList)
    vcl = len(vertList[0])
    ZERO_CUT = 1.0e-10
    print vrl, vcl
    BoolRes = numpy.zeros(vcl,'bool')
    INITIAL_ROW = vertList[0].copy()

    cntr = 0
    for r in xrange(vrl):
        BoolDiff = (numpy.absolute(vertList[r] - INITIAL_ROW) >= ZERO_CUT)
        for B in range(len(BoolDiff)):
            if BoolDiff[B]:
                BoolRes[B] = True
        if cntr == 1000:
            print('Processing vertex %s of %s' % (r, vrl))
            cntr = 0
        else:
            cntr += 1

    VarFluxes = []
    for f in range(len(BoolRes)):
        if BoolRes[f]:
            VarFluxes.append(FLUX_NAMES[f])

    print 'Variable fluxes:', len(VarFluxes)
    print VarFluxes
    CBTools.exportLabelledLinkedList([VarFluxes], None, os.path.join(work_dir,'%s.vertex_variable_fluxes.csv' % model_name), sep=',', format='%s', appendlist=False)
    CBTools.exportLabelledLinkedList([list(FLUX_NAMES), BoolRes.tolist()], None, os.path.join(work_dir,'%s.vertex_variable_fluxes_all.csv' % model_name), sep=',', format='%s', appendlist=False)


if WRITE_VERTEX_ARRAY_NEW:
    fname, VarVertArr = writeVertexArray_new((0,len(vertex_arr_status)), model_name, mod, vertList, VarFluxes)
    print 'VariableFluxVertices:', len(VarFluxes)

TIME_END = time.time()

print '\n\n'
print 'Time taken to analyse %s data: %2.2f min' % (model_name, (TIME_END-TIME_START)/60.0)
